package org.bdware.registry.handler;

import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import org.apache.log4j.Logger;
import org.bdware.doip.core.doipMessage.DoipMessage;
import org.bdware.doip.core.doipMessage.DoipResponseCode;
import org.bdware.doip.core.model.digitalObject.DigitalObject;
import org.bdware.doip.core.model.metadata.MetaDO;
import org.bdware.doip.core.model.metadata.SearchParameter;
import org.bdware.doip.core.model.metadata.SearchResult;
import org.bdware.doip.core.model.operations.BasicOperations;
import org.bdware.doip.core.utils.DoipGson;
import org.bdware.doip.endpoint.doipClient.DoipClientImpl;
import org.bdware.doip.endpoint.doipClient.DoipMessageCallback;
import org.bdware.doip.endpoint.doipServer.DoipServiceInfo;
import org.bdware.doip.endpoint.doipServer.Op;
import org.bdware.doip.endpoint.doipServer.RegistryHandlerBase;
import org.bdware.doip.implementation.DoIndexer;
import org.bdware.registry.storage.BDRegStorage;
import org.bdware.registry.storage.RemoteReg;
import org.rocksdb.RocksDBException;

import java.util.HashMap;

public class BDRegistryHandler extends RegistryHandlerBase {

    static Logger logger = Logger.getLogger(BDRegistryHandler.class);
    BDRegStorage store;
    DoIndexer indexer;
    HashMap<String, String> inRegs;
    HashMap<String, RemoteReg> outRegs;
    static final String IN_REGS = "REMOTEREG";
    static final String OUT_REGS = "outRegs";
    static final String DO_OWNER = "doOwner";

    public BDRegistryHandler(DoipServiceInfo info, BDRegStorage s, DoIndexer i) throws RocksDBException {
        super(info);
        this.serviceInfo = info;
        store = s;
        indexer = i;
        if(store.getRaw(IN_REGS.getBytes()) != null){
            inRegs = loadRegFederation();
        }else inRegs = new HashMap<>();
        if(store.getRaw(OUT_REGS.getBytes()) != null){
            outRegs = loadOutRegs();
        }else
            outRegs = new HashMap<>();
    }

    @Override
    public DoipMessage handleHello(DoipMessage request) {
        return replyDoipServiceInfo(request);
    }

    @Override
    public DoipMessage handleListOps(DoipMessage request) {
        return replyAllOperations(request);
    }

    @Override
    public DoipMessage handleCreate(DoipMessage request) {
        try {
            DigitalObject digitalObject = request.body.getDataAsDigitalObject();
            MetaDO metaDO = MetaDO.fromDO(digitalObject);
            if (metaDO == null || metaDO.id == null || metaDO.id.equals(""))
                return replyStringWithStatus(request, "Invalid metaDO.", DoipResponseCode.Invalid);
            if(getDoOwner(digitalObject)!= null )
                metaDO.addAttribute(DO_OWNER,getDoOwner(digitalObject));
            DigitalObject ret = store.getDoByID(metaDO.id);
            if (ret != null) {
                return handleUpdate(request);
            }
            String owner = getDoOwner(metaDO);
            logger.debug("Message body length: " + request.body.getLength());
            logger.info("create metada of DO: " + metaDO.id);
            store.setRequester(owner).storeDo(metaDO);
            indexer.indexDigitalObject(metaDO);
            return replyDO(request,metaDO);
        } catch (Exception e) {
            e.printStackTrace();
            return replyStringWithStatus(request, "Server Exception: " + e.getMessage(), DoipResponseCode.MoreThanOneErrors);
        }
    }

    @Override
    public DoipMessage handleUpdate(DoipMessage request) {
        try {
            DigitalObject digitalObject = request.body.getDataAsDigitalObject();
            MetaDO metaDO = MetaDO.fromDO(digitalObject);
            if (metaDO == null)
                return replyStringWithStatus(request, "Metadata field not found.", DoipResponseCode.Invalid);
            metaDO.addAttribute(DO_OWNER,getDoOwner(digitalObject));
            if (!request.header.parameters.id.equals(metaDO.id)) {
                return replyStringWithStatus(request, "Invalid request: different identifiers in DO and message header.", DoipResponseCode.Invalid);
            }
            DigitalObject ret = store.getDoByID(metaDO.id);
            if (ret == null) {
                return replyStringWithStatus(request, "Failed: DO is not existed.", DoipResponseCode.DoNotFound);
            }
            if (checkUpdateAuth(metaDO)) return replyStringWithStatus(request,"Failed: unAuthOperation", DoipResponseCode.UnAuth_Op);

            String owner = getDoOwner(metaDO);
            store.setRequester(owner).updateDo(metaDO);
            indexer.deleteDigitalObject(metaDO.id);
            indexer.indexDigitalObject(metaDO);
            return replyDO(request,metaDO);
        } catch (Exception e) {
            e.printStackTrace();
            return replyStringWithStatus(request, "Server Exception: " + e.getMessage(), DoipResponseCode.MoreThanOneErrors);
        }
    }

    @Override
    public DoipMessage handleDelete(DoipMessage request) {
        try {
            DigitalObject ret = store.getDoByID(request.header.parameters.id);
            if (ret == null) {
                return replyStringWithStatus(request, "Failed: DO is not existed.", DoipResponseCode.DoNotFound);
            }
            String owner = getDoOwner(ret);
            if(owner != null && !owner.equals(request.credential.getSigner()))
                return replyStringWithStatus(request,"Failed: unAuthOperation", DoipResponseCode.UnAuth_Op);
            store.setRequester(owner).deleteDoByID(request.header.parameters.id);
            indexer.deleteDigitalObject(request.header.parameters.id);
            return replyNull(request);
        }catch (Exception e){
            e.printStackTrace();
            return replyStringWithStatus(request, "Server Exception: " + e.getMessage(), DoipResponseCode.MoreThanOneErrors);
        }
    }

    @Override
    public DoipMessage handleRetrieve(DoipMessage request) {
        try {
            DigitalObject ret = store.getDoByID(request.header.parameters.id);
            if (ret == null) {
                return replyStringWithStatus(request, "Failed: DO is not existed.", DoipResponseCode.DoNotFound);
            }
            return replyDO(request,ret);
        }catch (Exception e){
            e.printStackTrace();
            return replyStringWithStatus(request, "Server Exception: " + e.getMessage(), DoipResponseCode.MoreThanOneErrors);
        }
    }

    @Override
    public DoipMessage handleUnknown(DoipMessage request) {
        return replyStringWithStatus(request,"Unsupported Operation!", DoipResponseCode.Declined);
    }

    @Override
    public DoipMessage handleSearch(DoipMessage request) {
        if(request.header.parameters.attributes == null ){
            logger.info("invalid search request: attributes not found");
            return replyStringWithStatus(request, "invalid search request: attributes not found.", DoipResponseCode.Invalid);
        }
        SearchParameter sp = new SearchParameter(request.header.parameters.attributes.get("query").getAsString(),
                !request.header.parameters.attributes.has("pageNum")?0:request.header.parameters.attributes.get("pageNum").getAsInt(),
                !request.header.parameters.attributes.has("pageSize")?0:request.header.parameters.attributes.get("pageSize").getAsInt(),
                !request.header.parameters.attributes.has("sortFields")?"":request.header.parameters.attributes.get("sortFields").getAsString(),
                !request.header.parameters.attributes.has("type")?null:request.header.parameters.attributes.get("type").getAsString());
        try {
            SearchResult sr = indexer.search(sp);
            distributedSearch(sr,request);
            return replyString(request, DoipGson.getDoipGson().toJson(sr));
        } catch (Exception e) {
            e.printStackTrace();
            return replyStringWithStatus(request, "Server Exception: " + e.getMessage(), DoipResponseCode.MoreThanOneErrors);
        }
    }

    @Op(op = BasicOperations.Extension,name = "Op.Join")
    public DoipMessage handleJoin(DoipMessage request){
        logger.info(new Gson().toJson(request));

        if(request.header.parameters.attributes == null ||request.header.parameters.attributes.get("repoID")== null){
            logger.info("invalid join request: repoID not found");
            return replyStringWithStatus(request, "invalid join request: repoID not found.", DoipResponseCode.Invalid);
        }
        String repoID = request.header.parameters.attributes.get("repoID").getAsString();
        String repoUrl = request.header.parameters.attributes.get("repoUrl").getAsString();
        logger.info(repoID);
        inRegs.put(repoID,repoUrl);
        saveRegFederation();
        return replyString(request,"success");
    }

    @Op(op = BasicOperations.Extension,name = "Op.Quit")
    public DoipMessage handleQuit(DoipMessage request){
        if(request.header.parameters.attributes == null ||request.header.parameters.attributes.get("repoID")== null){
            logger.info("invalid join request: repoID not found");
            return replyStringWithStatus(request, "invalid join request: repoID not found.", DoipResponseCode.Invalid);
        }
        String repoID = request.header.parameters.attributes.get("repoID").getAsString();
        inRegs.remove(repoID);
        saveRegFederation();
        return replyString(request,"success");
    }

    private boolean checkUpdateAuth(DigitalObject newDO) {
        DigitalObject oldDO = store.getDoByID(newDO.id);
        if(getDoOwner(oldDO) == null) return true;
        return getDoOwner(oldDO).equals(getDoOwner(newDO));
    }

    String getDoOwner(DigitalObject digitalObject){
        if(digitalObject.attributes.get(DO_OWNER) != null){
            logger.debug(digitalObject.attributes.get(DO_OWNER).getAsString());
            return digitalObject.attributes.get(DO_OWNER).getAsString();
        }
        return null;
    }

    void saveRegFederation(){
        try {
            store.putRaw(IN_REGS.getBytes(),new Gson().toJson(inRegs).getBytes());
            store.putRaw(OUT_REGS.getBytes(),new Gson().toJson(outRegs).getBytes());
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    HashMap<String,RemoteReg> loadOutRegs(){
        try {
            byte[] regBytes = store.getRaw(OUT_REGS.getBytes());
            return new Gson().fromJson(new String(regBytes),new TypeToken<HashMap<String,RemoteReg>>(){}.getType());
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
        return null;
    }

    HashMap<String,String> loadRegFederation(){
        try {
            byte[] regBytes = store.getRaw(IN_REGS.getBytes());
            return new Gson().fromJson(new String(regBytes),new TypeToken<HashMap<String,String>>(){}.getType());
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
        return null;
    }

    public HashMap<String, String> getInRegs(){
        return inRegs;
    }

    public HashMap<String,RemoteReg> getOutRegs() { return outRegs; }

    public void addOutRegs(RemoteReg newReg){
        outRegs.put(newReg.id,newReg);
        saveRegFederation();
    }

    void distributedSearch(SearchResult sr, DoipMessage searchMsg){
        AggregateFlag aggregateFlag = new AggregateFlag(inRegs.size());

        if(aggregateFlag.isFinished()) return;

        for (String url: inRegs.values()) {
            new Thread(()->doSearch(url,searchMsg, new searchCallback(sr, aggregateFlag))).start();
        }
        int wait = 4;           //wait result for 2s
        while(!aggregateFlag.isFinished() && wait>0){
            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            wait--;
        }
        aggregateFlag.abort();
    }

    void doSearch(String repoURL, DoipMessage searchMsg, DoipMessageCallback cb){

        DoipClientImpl client = new DoipClientImpl();
        client.connect(repoURL);
        logger.debug("send search message to :" + repoURL);
        client.sendRawMessage(searchMsg,cb);

    }

    static class searchCallback implements DoipMessageCallback{

        final SearchResult sr;
        final AggregateFlag wf;

        public searchCallback(SearchResult sr, AggregateFlag flag){
            this.sr = sr;
            wf = flag;
        }

        @Override
        public void onResult(DoipMessage msg) {
            logger.debug("receive a search resp: " + msg.body.getDataAsJsonString());
            if(wf.isFinished()) return;
            if(msg.header.parameters.response != DoipResponseCode.Success){
                logger.warn("search failed, target: " + msg.header.parameters.id + "; error msg: " + msg.body.getDataAsJsonString());
                return;
            }
            SearchResult result;
            result = new Gson().fromJson(msg.body.getDataAsJsonString(),SearchResult.class);
            synchronized (sr){
                sr.results.addAll(result.results);
                sr.size += result.size;
            }
            synchronized (wf){
                wf.receive();
            }
        }
    }

    static class AggregateFlag {
        int received = 0;
        final int total;
        public AggregateFlag(int total){
            this.total = total;
        }
        public void receive(){ received ++;}
        public boolean isFinished(){
            return received >= total;
        }
        void abort(){
            received = total;
        }
    }
}
